#include "lib/print.h"
#include "lib/str.h"
#include "mem/pmem.h"
#include "mem/vmem.h"
#include "mem/mmap.h"
#include "proc/cpu.h"
#include "proc/initcode.h"
#include "memlayout.h"
#include "riscv.h"
#include "fs/fs.h" // 
#include "proc/proc.h"  // 确保包含 proc_wakeup_one 声明的头文件

/*----------------外部空间------------------*/

// in trampoline.S
extern char trampoline[];

// in swtch.S
extern void swtch(context_t* old, context_t* new);

// in trap_user.c
extern void trap_user_return();

// defined in kvm.c
extern pgtbl_t kernel_pgtbl;

/*----------------本地变量------------------*/

// 进程数组
static proc_t procs[NPROC];

// 第一个进程的指针
static proc_t* proczero;

// 全局的pid和保护它的锁 
static int global_pid = 1;
static spinlock_t lk_pid;


// 申请一个pid(锁保护)
static int alloc_pid()
{
    int tmp = 0;
    spinlock_acquire(&lk_pid);
    assert(global_pid >= 0, "alloc_pid: overflow");
    tmp = global_pid++;
    spinlock_release(&lk_pid);
    return tmp;
}

// 释放锁 + 调用 trap_user_return
static void fork_return()
{
    proc_t* p = myproc();
    spinlock_release(&p->lk);

    if (p == proczero) {
        // 初始化
        fs_init();
    }

    trap_user_return();
}

// 返回一个未使用的进程空间
// 设置pid + 设置上下文中的ra和sp
// 申请tf和pgtbl使用的物理页
proc_t* proc_alloc() {
    proc_t *p;

    // 遍历进程数组，查找状态为 UNUSED 的进程
    for (p = &procs[0]; p < &procs[NPROC]; p++) {
        spinlock_acquire(&p->lk);

        if (p->state == UNUSED) {
            // 设置进程相关字段
            p->pid = alloc_pid(); // 分配一个新的 PID

            // 分配 trapframe 和页表使用的物理页
            p->tf = pmem_alloc(false);
            if (!p->tf) {
                proc_free(p); // 回收资源
                spinlock_release(&p->lk);
                return NULL; // 分配失败
            }

            p->pgtbl = proc_pgtbl_init((uint64)p->tf);
            if (!p->pgtbl) {
                proc_free(p); // 回收资源
                spinlock_release(&p->lk);
                return NULL; // 分配失败
            }

            // 设置上下文中的 sp 和 ra
            memset(&p->ctx, 0, sizeof(p->ctx));
            p->ctx.ra = (uint64)fork_return;        // 返回地址
            p->ctx.sp = p->kstack + PGSIZE;         // 内核栈指针

            // 返回已分配的进程，锁未释放
            return p;
        }

        // 如果当前进程不符合条件，释放锁
        spinlock_release(&p->lk);
    }

    return NULL; // 没有找到可用进程
}

// 释放一个进程空间
// 释放pgtbl的整个地址空间
// 释放mmap_region到仓库
// 设置其余各个字段为合适初始值
// tips: 调用者需持有p->lk
void proc_free(proc_t* proc) {
    // 释放 trapframe
    if (proc->tf) {
        pmem_free((uint64)proc->tf, false);
        proc->tf = NULL; // 防止悬空指针
    }

    // 销毁页表
    if (proc->pgtbl) {
        uvm_destroy_pgtbl(proc->pgtbl); // 销毁进程页表
        proc->pgtbl = NULL;            // 防止悬空指针
    }

    // 释放 mmap_region
    mmap_region_t* current_region = proc->mmap;
    while (current_region) {
        mmap_region_t* temp = current_region;
        current_region = current_region->next;
        mmap_region_free(temp); // 释放单个映射区域的内存
    }
    proc->mmap = NULL; // 重置 mmap 指针

    // 重置进程字段
    proc->pid = -1;
    proc->state = UNUSED;
    proc->parent = NULL;
    proc->exit_state = 0;

    // 清空进程内核上下文
    memset(&(proc->ctx), 0, sizeof(context_t)); // 重置内核上下文结构
}

// 进程模块初始化
void proc_init() {
    proc_t *proc;

    // 初始化全局 PID 锁
    spinlock_init(&lk_pid, "global_pid_lock");

    // 初始化 proczero 指针，指向进程数组第一个进程
    proczero = &procs[0];

    // 遍历所有进程
    for (proc = &procs[0]; proc < &procs[NPROC]; proc++) {
        // 初始化每个进程的自旋锁
        spinlock_init(&proc->lk, "proc_lock");

        // 为进程分配物理页作为内核栈
        void *kstack_page = pmem_alloc(true); // 分配物理页并清零
        if (kstack_page == 0) {
            panic("proc_init: pmem_alloc failed.");
        }

        // 获取内核栈的虚拟地址
        uint64 kernel_stack_addr = KSTACK((int)(proc - &procs[0]));

        // 将分配的物理内存映射到内核页表
        vm_mappages(kernel_pgtbl, kernel_stack_addr, (uint64)kstack_page, PGSIZE, PTE_R | PTE_W);

        // 保存内核栈虚拟地址
        proc->kstack = kernel_stack_addr;

        // 初始化进程的其他字段
        spinlock_acquire(&proc->lk);
        proc_free(proc); // 将进程标记为空闲并释放资源
        spinlock_release(&proc->lk);
    }

    // 刷新内核页表，确保 TLB 生效
    kvm_inithart();
}


// 获得一个初始化过的用户页表
// 完成了trapframe 和 trampoline 的映射
pgtbl_t proc_pgtbl_init(uint64 trapframe)
{
    // 分配页表
    pgtbl_t pagetable = (pgtbl_t)pmem_alloc(true);

    // memset(pagetable, 0, PGSIZE);
    if (pagetable == 0)
        panic("proc_pgtbl_init: failed!\n");

    // 映射trampoline
    vm_mappages(pagetable, TRAMPOLINE, (uint64)trampoline, PGSIZE, PTE_R | PTE_X);
    // 映射trapframe
    vm_mappages(pagetable, TRAPFRAME, trapframe, PGSIZE, PTE_R | PTE_W);
    
    return pagetable;
}

/*
    第一个用户态进程的创建
    它的代码和数据位于initcode.h的initcode数组

    第一个进程的用户地址空间布局:
    trapoline   (1 page)
    trapframe   (1 page)
    ustack      (1 page)
    .......
                        <--heap_top
    code + data (1 page)
    empty space (1 page) 最低的4096字节 不分配物理页，同时不可访问
*/
void proc_make_first() {

    uint64 page;

    printf("Debug: Entering proc_make_first\n");
    printf("Debug: proczero address = ");
    print_hex((uint64)proczero);

    if (!proczero) {
        panic("proc_make_first: proczero is not initialized!");
    }

    // pid 设置
    proczero->pid = 1;

    // pagetable 初始化
    proczero->tf = (trapframe_t *)pmem_alloc(false);
    if (!proczero->tf) {
        panic("Failed to allocate trapframe.");
    } else {
        printf("Debug: Trapframe allocated at ");
        print_hex((uint64)proczero->tf);
    }

    proczero->pgtbl = proc_pgtbl_init((uint64)proczero->tf);
    if (!proczero->pgtbl) {
        panic("proc_make_first: Page table not initialized.");
    } else {
        printf("Debug: Page table initialized at address = ");
        print_hex((uint64)proczero->pgtbl);
    }

    // ustack 映射 + 设置 ustack_pages 
    char *ustack_page = pmem_alloc(false);  // 分配用户栈物理页
    if (!ustack_page) {
        panic("proc_make_first: Failed to allocate user stack.");
    } else {
        printf("Debug: User stack allocated at physical address = ");
        print_hex((uint64)ustack_page);
    }

    vm_mappages(proczero->pgtbl, TRAPFRAME - PGSIZE, (uint64)ustack_page, PGSIZE, PTE_R | PTE_W | PTE_U);
    // printf("Debug: User stack mapped at virtual address = ");
    // printf("Virtual Address: 0x%lx\n", (uint64)(TRAPFRAME - PGSIZE));

    proczero->ustack_pages = 1;  // 用户栈页面数为 1

    // data + code 映射
    assert(initcode_len <= PGSIZE, "proc_make_first: initcode too big\n");

    page = (uint64)pmem_alloc(false);  // 分配一个物理页
    if (!page) {
        panic("proc_make_first: Failed to allocate code page.");
    } else {
        printf("Debug: Code page allocated at physical address = ");
        print_hex(page);
    }

    memset((void*)page, 0, PGSIZE);  // 清空物理页内容
    memmove((void*)page, initcode, initcode_len);  // 复制 initcode 数据

    vm_mappages(proczero->pgtbl, USER_BASE, page, PGSIZE, PTE_W | PTE_R | PTE_X | PTE_U);
    printf("Debug: User code mapped at virtual address = ");
    print_hex(USER_BASE);

    // 设置 heap_top
    proczero->heap_top = USER_BASE + PGSIZE;  // 初始化 heap_top 到用户基地址 + 代码段大小

    // tf字段设置
    proczero->tf->epc = USER_BASE;  // 设置程序计数器为代码段起始地址
    proczero->tf->sp = TRAPFRAME;  // 设置用户栈指针为栈底地址
    printf("Debug: User epc set to ");
    print_hex(proczero->tf->epc);
    printf("Debug: User sp set to ");
    print_hex(proczero->tf->sp);

    // 内核字段设置
    proczero->kstack = KSTACK(proczero->pid);  // 分配内核栈
    proczero->ctx.sp = proczero->kstack + PGSIZE;  // 内核栈指针指向栈顶
    proczero->ctx.ra = (uint64)trap_user_return;  // 返回地址设为 trap_user_return
    printf("Debug: Kernel stack allocated at virtual address = ");
    print_hex(proczero->kstack);
    printf("Debug: Kernel context sp = ");
    print_hex(proczero->ctx.sp);
    printf(", ra = ");
    print_hex(proczero->ctx.ra);

    // 设置 mmap_region_t
    proczero->mmap = mmap_region_alloc();
    if (!proczero->mmap) {
        panic("proc_make_first: Failed to allocate initial mmap_region_t.");
    }

    proczero->mmap->begin = MMAP_BEGIN; // 起始地址
    proczero->mmap->npages = (MMAP_END - MMAP_BEGIN) / PGSIZE; // 管理的页面数量
    proczero->mmap->next = NULL; // 链表结束
    printf("Debug: mmap region initialized from ");
    print_hex(proczero->mmap->begin);
    printf(" to ");
    print_hex(proczero->mmap->begin + proczero->mmap->npages * PGSIZE);

    // 上下文切换
    cpu_t *c = mycpu();  // 获取当前 CPU
    c->proc = proczero;  // 设置当前进程为 proczero
    printf("Debug: About to switch to proczero context\n");
    swtch(&c->ctx, &proczero->ctx);
    printf("Debug: Back from swtch (this should not happen!)\n");
}



// 进程复制
// UNUSED -> RUNNABLE
int proc_fork() {
    proc_t* parent = myproc(); // 获取当前进程（父进程）

    // 分配一个新的子进程
    proc_t* child = proc_alloc();
    if (child == NULL) {
        return -1; // 分配失败，返回错误码
    }

    // 加锁保护子进程，防止在初始化过程中被调度
    spinlock_acquire(&child->lk);

    // 复制父进程的 trapframe
    *(child->tf) = *(parent->tf);
    child->tf->a0 = 0; // 子进程的返回值为 0

    // 复制父进程的页表
    uvm_copy_pgtbl(parent->pgtbl, child->pgtbl, parent->heap_top, parent->ustack_pages, parent->mmap);

    // 复制 mmap 链表
    if (parent->mmap) {
        child->mmap = mmap_region_alloc();
        if (child->mmap == NULL) {
            // mmap 分配失败，释放资源并返回错误码
            proc_free(child);
            spinlock_release(&child->lk);
            return -1;
        }

        mmap_region_t* src_region = parent->mmap;
        mmap_region_t* dst_region = child->mmap;

        while (src_region) {
            dst_region->begin = src_region->begin;
            dst_region->npages = src_region->npages;

            if (src_region->next) {
                dst_region->next = mmap_region_alloc();
                if (dst_region->next == NULL) {
                    // 如果分配失败，释放子进程的资源
                    proc_free(child);
                    spinlock_release(&child->lk);
                    return -1;
                }
            } else {
                dst_region->next = NULL;
            }

            src_region = src_region->next;
            dst_region = dst_region->next;
        }
    } else {
        child->mmap = NULL;
    }

    // 复制堆顶和用户栈页数
    child->heap_top = parent->heap_top;
    child->ustack_pages = parent->ustack_pages;

    // 设置子进程的父进程
    child->parent = parent;

    // 设置子进程为 RUNNABLE，表示可以被调度
    child->state = RUNNABLE;

    // 释放子进程的锁
    spinlock_release(&child->lk);

    // 返回子进程的 PID
    return child->pid;
}

// 进程放弃CPU的控制权
// RUNNING -> RUNNABLE
void proc_yield()
{
    proc_t *proc = myproc();

    spinlock_acquire(&proc->lk);

    proc->state = RUNNABLE;
    //转到调度器
    proc_sched();

    spinlock_release(&proc->lk);
}


// 等待一个子进程进入 ZOMBIE 状态
// 将退出的子进程的exit_state放入用户给的地址 addr
// 成功返回子进程pid，失败返回-1
int proc_wait(uint64 addr) {
    proc_t *np;             // 用于遍历子进程
    int have_child = 0;     // 标记是否存在子进程
    int pid = -1;           // 返回的子进程 pid
    proc_t *parent = myproc(); // 当前进程，即父进程

    spinlock_acquire(&parent->lk); // 获取父进程的锁

    while (1) {
        have_child = 0; // 初始化标记

        // 遍历进程表
        for (np = procs; np < &procs[NPROC]; np++) {
            if (np->parent == parent) { // 找到子进程
                have_child = 1;

                spinlock_acquire(&np->lk); // 获取子进程的锁

                if (np->state == ZOMBIE) { // 找到退出的子进程
                    pid = np->pid;

                // 将子进程的退出状态复制到用户指定地址
                    uvm_copyout(parent->pgtbl, addr, (uint64)&np->exit_state, sizeof(np->exit_state));
                    // if (uvm_copyout(parent->pgtbl, addr, (uint64)&np->exit_state, sizeof(np->exit_state)) != 0) {
                    //     spinlock_release(&np->lk);
                    //     spinlock_release(&parent->lk);
                    //     return -1; // 用户地址非法
                    // }

                    // 清理子进程的资源
                    proc_free(np);

                    // 释放锁
                    spinlock_release(&np->lk);
                    spinlock_release(&parent->lk);

                    return pid; // 返回子进程的 pid
                }

                // 释放子进程的锁
                spinlock_release(&np->lk);
            }
        }

        // 如果没有子进程，返回 -1
        if (!have_child) {
            spinlock_release(&parent->lk);
            return -1;
        }

        // 如果有子进程但没有处于 ZOMBIE 状态，则等待
        proc_sleep(parent, &parent->lk); // 父进程进入睡眠等待子进程状态更新
    }
}

// 父进程退出，子进程认proczero做父，因为它永不退出
static void proc_reparent(proc_t* parent)
{
    proc_t *p;
    for(p = &procs[0]; p < &procs[NPROC]; p++)
    {
        if(p->parent == parent)
        {
            p->parent = proczero;
            // 唤醒proczero，使其及时检测到子进程的退出并释放子进程的空间
            proc_wakeup_one(proczero);
        }
    }
}


// 唤醒一个进程
void proc_wakeup_one(proc_t* p)
{
    assert(spinlock_holding(&p->lk), "proc_wakeup_one: lock");
    if(p->state == SLEEPING && p->sleep_space == p) {
        p->state = RUNNABLE;
    }
}

// 进程退出
void proc_exit(int exit_state) {
    proc_t *p = myproc(); // 获取当前退出的进程

    // 1号进程 (proczero) 不能退出
    if (p == proczero) {
        panic("proc_exit: proczero can not exit");
    }

    // 唤醒 proczero（防止孤儿进程积累）
    spinlock_acquire(&proczero->lk);
    proc_wakeup_one(proczero);
    spinlock_release(&proczero->lk);

    // 获取父进程的副本（避免父进程状态变化造成的影响）
    spinlock_acquire(&p->lk);
    proc_t *parent = p->parent; // 获取父进程的引用
    spinlock_release(&p->lk);

    // 将子进程交给 proczero 监管（孤儿进程处理）
    spinlock_acquire(&p->lk);
    proc_reparent(p); // 将当前进程的子进程的 parent 指针重定向为 proczero
    spinlock_release(&p->lk);

    // 唤醒父进程（用于清理当前退出进程的资源）
    spinlock_acquire(&parent->lk);
    proc_wakeup_one(parent); // 唤醒等待该进程退出的父进程
    spinlock_release(&parent->lk);

    // 设置当前进程的退出状态，并将其状态置为 ZOMBIE
    spinlock_acquire(&p->lk);
    p->exit_state = exit_state; // 设置退出状态
    p->state = ZOMBIE;          // 将状态置为 ZOMBIE
    spinlock_release(&p->lk);

    // 切换到调度器，运行其他可运行的进程
    proc_sched();

    // 如果调度器没有运行，触发 panic（理论上不会到达这里）
    panic("proc_exit: zombie exit");
}

// 进程切换到调度器
// ps: 调用者保证持有当前进程的锁
void proc_sched() {
    proc_t *p = myproc(); // 获取当前进程指针

    // 确保调用者持有当前进程的锁
    if (!spinlock_holding(&p->lk)) {
        printf("proc_sched: lock not held by process %d\n", p->pid);
        panic("proc_sched: lock not held.");
    }

    // 确保当前 CPU 的 `noff` 为 1（内核线程锁计数为 1）
    if (mycpu()->noff != 1) {
        printf("proc_sched: CPU noff = %d\n", mycpu()->noff);
        panic("proc_sched: CPU noff != 1.");
    }

    // 确保当前进程不是 `RUNNING` 状态
    if (p->state == RUNNING) {
        printf("proc_sched: process %d is still running\n", p->pid);
        panic("proc_sched: running process cannot schedule.");
    }

    // 确保设备中断已禁用
    if (intr_get()) {
        panic("proc_sched: device interrupts enabled.");
    }

    // 保存当前 CPU 状态
    int origin = mycpu()->origin;

    // 通过上下文切换保存当前进程上下文并切换到 CPU 的上下文
    swtch(&p->ctx, &mycpu()->ctx);

    // 恢复保存的 CPU 状态
    mycpu()->origin = origin;
}

// 调度器
void proc_scheduler() {
    proc_t *p;              // 指向进程表中的某个进程
    cpu_t *c = mycpu();     // 当前 CPU 的信息

    // 设置当前 CPU 核心的运行进程为空
    c->proc = NULL;

    for (;;) {
        // 启用中断，允许设备中断发生
        intr_on();

        // 遍历进程表，寻找 `RUNNABLE` 状态的进程
        for (p = &procs[0]; p < &procs[NPROC]; p++) {
            spinlock_acquire(&p->lk); // 获取进程锁

            if (p->state == RUNNABLE) {
                // 进程状态切换为 `RUNNING`
                printf("proc_scheduler: process %d switching to RUNNING\n", p->pid); // 调试信息
                p->state = RUNNING;

                // 让当前 CPU 执行此进程
                c->proc = p;
                swtch(&c->ctx, &p->ctx); // 上下文切换到选定进程

                // 进程切换回来时，表示进程已暂停或退出，重置当前运行进程
                c->proc = NULL;
            }

            spinlock_release(&p->lk); // 释放进程锁
        }
    }
}

// 进程睡眠在sleep_space
void proc_sleep(void* sleep_space, spinlock_t* lk) {
    proc_t *p = myproc(); // 获取当前进程

    // 确保传入的锁不为空，避免潜在错误
    assert(lk != NULL, "proc_sleep: lock is NULL");

    // 如果传入的锁不是当前进程的锁，切换锁
    if (lk != &p->lk) {
        spinlock_acquire(&p->lk); // 获取当前进程的锁，确保同步状态
        spinlock_release(lk);     // 释放传入的锁，防止死锁
    }

    // 设置进程为睡眠状态
    p->sleep_space = sleep_space; // 设置睡眠原因标记（通常是等待的资源地址）
    p->state = SLEEPING;          // 将进程状态设置为 SLEEPING

    // 让出 CPU，切换到调度器
    proc_sched();

    // 当进程被唤醒后恢复执行，取消睡眠事件
    p->sleep_space = NULL; // 清空睡眠原因标记

    // 如果原始锁和当前进程锁不同，切换回原始锁
    if (lk != &p->lk) {
        spinlock_release(&p->lk); // 释放当前进程的锁
        spinlock_acquire(lk);     // 恢复原始锁
    }
}

// 唤醒所有在sleep_space沉睡的进程
void proc_wakeup(void* sleep_space)
{
    proc_t *proc;
    for(proc = &procs[0]; proc < &procs[NPROC]; proc++) 
    {
        spinlock_acquire(&proc->lk);
        if(proc->state == SLEEPING && proc->sleep_space == sleep_space)
            proc->state = RUNNABLE;
        spinlock_release(&proc->lk);
    }
}
